package com.heap;

import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;

public class TopK {
    public static void main(String[] args) {
        int[] nums = {3, 13, 6, 9, 10};
        int k = 3;
        int[] result = topK(nums, k);
        System.out.println(Arrays.toString(result));
    }

    /**
     * 最大堆
     * @param nums
     * @param k
     * @return
     */
    public static int[] topK2(int[] nums, int k){
        int[] result = new int[k];
        if (nums == null || nums.length == 0) return result;

        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });

        for (int num : nums) {
            maxHeap.add(num);
        }

        for (int i = 0; i < result.length; i++) {
            result[i] = maxHeap.poll();
        }
        return result;
    }

    /**
     * 最小堆
     * @param nums
     * @param k
     * @return
     */
    public static int[] topK(int[] nums, int k) {
        int[] result = new int[k];
        if (nums == null || nums.length == 0) return result;

        PriorityQueue<Integer> minHeap = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o1 - o2;
            }
        });

        for (int num : nums) {
            minHeap.add(num);
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }

        for (int i = 0; i < result.length; i++) {
            result[k - i - 1] = minHeap.poll();
        }
        return result;
    }
}
